package com.hillfly.mybatis.shard.plugin;

import com.hillfly.mybatis.shard.ShardHelper;
import com.hillfly.mybatis.shard.converter.SqlConverterFactory;
import com.hillfly.mybatis.shard.utils.Reflections;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.executor.statement.BaseStatementHandler;
import org.apache.ibatis.executor.statement.RoutingStatementHandler;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.util.Properties;

@Intercepts({@Signature(
        type = StatementHandler.class,
        method = "prepare",
        args = {Connection.class, Integer.class}
)})
public class MyBatisShardInterceptor implements Interceptor {

    private Logger logger = LoggerFactory.getLogger(getClass());

    public static final String DELEGATE = "delegate";
    public static final String MAPPED_STATEMENT = "mappedStatement";

    @Override
    public Object intercept(Invocation ivk) throws Throwable {
        try {
            if (ivk.getTarget().getClass().isAssignableFrom(RoutingStatementHandler.class)) {
                final RoutingStatementHandler statementHandler = (RoutingStatementHandler) ivk.getTarget();
                final BaseStatementHandler delegate = (BaseStatementHandler) Reflections.getFieldValue(statementHandler, DELEGATE);
                final MappedStatement mappedStatement = (MappedStatement) Reflections.getFieldValue(delegate, MAPPED_STATEMENT);
                final BoundSql boundSql = statementHandler.getBoundSql();
                String sql = SqlConverterFactory.getGlobalInstance().convert(boundSql, mappedStatement);
                if (StringUtils.isNotBlank(sql)) {
                    logger.debug("[Modified sql] " + sql);
                    Reflections.setFieldValue(statementHandler.getBoundSql(), "sql", sql);
                }
            }
            return ivk.proceed();
        } finally {
            if (ShardHelper.isShardParamsEmpty()) {
                ShardHelper.clearShardParams();
            }
        }
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
    }
}
